# -*-encoding: Utf-8 -*-
# @File : get_cookies.py
# author: HJT
# Time：2023/11/13
import requests
import warnings
import re

from commons.encrypt import encrypt_password
from config import settings


def login_user_manage(base_url, user_home_url, username, password):
    """
    :param base_url: 统一登录接口，如：http://xx.com/login/
    :param user_home_url: 用户管理saas的任一url
    :param username: 登录用户名
    :param password: 密码
    :return: cookie实例
    """
    login_url = base_url + '/login/'
    warnings.filterwarnings('ignore')
    # 公共配置
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
                      "Chrome/115.0.0.0 Safari/537.36",
    }

    # 请求1：get请求蓝鲸登录页面
    session = requests.session()
    login_page = session.get(login_url, headers=headers, verify=False)

    # ①提取登录页面html中，name="csrfmiddlewaretoken"
    # ②提取加密公钥PASSWORD_RSA_PUBLIC_KEY
    set_cookie = login_page.headers['Set-Cookie']
    csrfmiddlewaretoken = re.search(r"bklogin_csrftoken=(.*?);", set_cookie).group(1)
    PASSWORD_RSA_PUBLIC_KEY = re.search(r'PASSWORD_RSA_PUBLIC_KEY = "(.*?)"', login_page.text).group(1)

    # 2. 请求2：蓝鲸控制台登录
    # 密码加密处理
    password = encrypt_password(password, PASSWORD_RSA_PUBLIC_KEY)
    # 登录参数
    login_data = {
        "csrfmiddlewaretoken": csrfmiddlewaretoken,  # data中的重要参数，服务器会校验
        "username": username,
        "password": password,
        "next": "",
        "app_id": ""
    }

    headers['Referer'] = login_url  # headers中的重要参数，服务器会校验
    login_resp = session.post(login_url, data=login_data, headers=headers, verify=False)
    # print(f'控制台页面：{login_resp.text}')
    assert login_resp.status_code < 400, f'{login_url}登录失败，请检查用户名、密码'
    assert '桌面 | 腾讯蓝鲸智云' in re.search(r"<title>(.*?)</title>", login_resp.text).group(
        1), f'请求失败，URL:{login_url}'

    # ===================== 以上登录至控制台通用； 以下为各类Saas的Cookie追加获取 ==============================
    # 请求3：访问用户管理模块，为了获取Saas自身定制的cookie
    user_page = session.get(user_home_url, headers=headers, verify=False)
    # print(f'用户管理页面：{user_page.text}')
    assert user_page.status_code < 300, f'请求失败，URL:{user_home_url}'
    assert '用户管理 | 腾讯蓝鲸智云' in re.search(r"<title>(.*?)</title>", user_page.text).group(
        1), f'请求失败，URL:{user_home_url}'

    # 获取用户身份信息(cookie)
    cookies = session.cookies.get_dict()
    if cookies is None:
        raise Exception("登录失败，cookie为空值")

    cookies_str = ''.join([f'{k}={v};' for k, v in cookies.items()])

    set_cookie_header = user_page.headers.get('Set-Cookie')
    token = re.search('bkuser_csrftoken=(.*?);', set_cookie_header).group(1)

    return cookies_str, token


if __name__ == '__main__':
    base_url = settings.cfg.get('base_url')
    user_home_url = settings.cfg.get('user_manager').get('user_base_url') + settings.cfg.get('user_manager').get(
        'home_route')
    username = settings.cfg.get('superman').get('username')
    password = settings.cfg.get('superman').get('password')
    print(login_user_manage(base_url, user_home_url, username, password))
